#include "inc/core/ucgrid3d.h"
#include "stdlib.h"
#include "stdio.h"
#include "math.h"
#include "gsl/gsl_randist.h"
#include "inc/core/lib.h"

void empty_grid(ucgrid3d* grid)
{
    size_t size = grid->Nx * grid->Ny * grid->Nz;
    memset(grid->rho, 0, size * sizeof (floatdig));
    memset(grid->phi, 0, size * sizeof (floatdig));
    memset(grid->j, 0, size * sizeof (current3d));
}

ucgrid3d* ucgrid3d_init_grid(int Nx, int Ny, int Nz,
                             floatdig dx, floatdig dy, floatdig dz)
{
    int i;
    int j;
    int k;
    ucgrid3d* grid;
#ifdef VERBOSE
    printf("ucgrid3d_init_grid: init grid with sizes Nx: %d, Ny: %d, Nz: %d\n", Nx, Ny, Nz);
#endif
    grid = (ucgrid3d*) malloc(sizeof (ucgrid3d));
    grid->dx = dx;
    grid->dy = dy;
    grid->dz = dz;
    grid->Nx = Nx;
    grid->Ny = Ny;
    grid->Nz = Nz;
    grid->rho = (floatdig***) malloc(Nx * sizeof (floatdig**));
    grid->phi = (floatdig***) malloc(Nx * sizeof (floatdig**));
    grid->j = (current3d***) malloc(Nx * sizeof (current3d**));
    for (i = 0; i < Nx; i++)
    {
        grid->rho[i] = (floatdig**) malloc(Ny * sizeof (floatdig*));
        grid->phi[i] = (floatdig**) malloc(Ny * sizeof (floatdig*));
        grid->j[i] = (current3d**) malloc(Ny * sizeof (current3d*));
        for (j = 0; j < Ny; j++)
        {
            grid->rho[i][j] = (floatdig*) malloc(Nz * sizeof (floatdig));
            grid->phi[i][j] = (floatdig*) malloc(Nz * sizeof (floatdig));
            grid->j[i][j] = (current3d*) malloc(Nz * sizeof (current3d));
            for (k = 0; k < Nz; k++)
            {
                grid->rho[i][j][k] = 0;
                grid->phi[i][j][k] = 0;
                grid->j[i][j][k].v[X] = 0;
                grid->j[i][j][k].v[Y] = 0;
                grid->j[i][j][k].v[Z] = 0;
            }
        }
    }
    return grid;
}

void ucgrid3d_free_grid(ucgrid3d* grid)
{
    int Nx = grid->Nx;
    int Ny = grid->Ny;
    int i;
    int j;

#ifdef VERBOSE
    printf("ucgrid3d_free_grid: Free memory from grid's data\n");
#endif
    for (i = 0; i < Nx; i++)
    {
        for (j = 0; j < Ny; j++)
        {
            free(grid->rho[i][j]);
            free(grid->phi[i][j]);
            free(grid->j[i][j]);
            grid->rho[i][j] = NULL;
            grid->phi[i][j] = NULL;
            grid->j[i][j] = NULL;


        }
        free(grid->rho[i]);
        free(grid->phi[i]);
        free(grid->j[i]);
        grid->rho[i] = NULL;
        grid->phi[i] = NULL;
        grid->j[i] = NULL;
    }
    free(grid);
    grid = NULL;

}

void ucgrid3d_get_cell_pos(ucgrid3d* grid, npos3d* n, pos3d* p)
{
    p->v[X] = n->v[X] * grid->dx;
    p->v[Y] = n->v[Y] * grid->dy;
    p->v[Z] = n->v[Z] * grid->dz;
}

void ucgrid3d_get_full_size_of_system(ucgrid3d* grid, pos3d* Ls)
{
    Ls->v[X] = grid->Nx * grid->dx;
    Ls->v[Y] = grid->Ny * grid->dy;
    Ls->v[Z] = grid->Nz * grid->dz;
}

void ucgrid3d_get_neighbour_cells(ucgrid3d* grid, pos3d* ppos, npos3d* lc, npos3d* rc)
{
    floatdig x = ppos->v[X] / grid->dx;
    floatdig y = ppos->v[Y] / grid->dy;
    floatdig z = ppos->v[Z] / grid->dz;

    lc->v[X] = floor(x);
    rc->v[X] = ceil(x);
    lc->v[Y] = floor(y);
    rc->v[Y] = ceil(y);
    lc->v[Z] = floor(z);
    rc->v[Z] = ceil(z);
}

void ucgrid3d_weight_particles_rho(particles_set* set, ucgrid3d* grid)
{
#ifdef VERBOSE
    printf("ucgrid3d_weight_particles_rho: weightening\n");
    printf("ucgrid3d_weight_particles_rho: using weightening method: %d\n",
           WEIGHTENING_METHOD, BOUNDARIES_TYPE);
#endif
#if ((WEIGHTENING_METHOD == WLINEAR) && (BOUNDARIES_TYPE == CYCLIC_BOUNDARIES))
    npos3d nlc = {0};
    npos3d nrc = {0};
    pos3d plc = {0};
    pos3d prc = {0};

    pos3d plll = {0};
    pos3d prll = {0};
    pos3d plrl = {0};
    pos3d plrr = {0};
    pos3d pllr = {0};
    pos3d prlr = {0};
    pos3d prrr = {0};
    pos3d prrl = {0};

    floatdig lllq = 0;
    floatdig rllq = 0;
    floatdig lrlq = 0;
    floatdig lrrq = 0;
    floatdig llrq = 0;
    floatdig rlrq = 0;
    floatdig rrrq = 0;
    floatdig rrlq = 0;

    pos3d* ppos;
    speed3d* pspeed;
    int i;
    floatdig dv = grid->dx * grid->dy * grid->dz;
    floatdig qdv = set->charge / dv;
    for (i = 0; i < set->num_of_particles; i++)
    {
        if (set->particles[i].active == PARTICLE_ACTIVE)
        {
            ppos = &(set->particles[i].pos);
            pspeed = &(set->particles[i].speed);
            ucgrid3d_get_neighbour_cells(grid, ppos, &nlc, &nrc);
            ucgrid3d_get_cell_pos(grid, &nlc, &plc);
            ucgrid3d_get_cell_pos(grid, &nrc, &prc);

            plll.v[X] = plc.v[X];
            plll.v[Y] = plc.v[Y];
            plll.v[Z] = plc.v[Z];

            prll.v[X] = prc.v[X];
            prll.v[Y] = plc.v[Y];
            prll.v[Z] = plc.v[Z];

            plrl.v[X] = plc.v[X];
            plrl.v[Y] = prc.v[Y];
            plrl.v[Z] = plc.v[Z];

            prrl.v[X] = prc.v[X];
            prrl.v[Y] = prc.v[Y];
            prrl.v[Z] = plc.v[Z];

            prrr.v[X] = prc.v[X];
            prrr.v[Y] = prc.v[Y];
            prrr.v[Z] = prc.v[Z];

            prlr.v[X] = prc.v[X];
            prlr.v[Y] = plc.v[Y];
            prlr.v[Z] = prc.v[Z];

            pllr.v[X] = plc.v[X];
            pllr.v[Y] = plc.v[Y];
            pllr.v[Z] = prc.v[Z];

            plrr.v[X] = plc.v[X];
            plrr.v[Y] = prc.v[Y];
            plrr.v[Z] = prc.v[Z];

            lllq = qdv * (prll.v[X] - ppos->v[X])*(plrl.v[Y] - ppos->v[Y])*(pllr.v[Z] - ppos->v[Z]) / dv;
            rllq = qdv * (ppos->v[X] - plll.v[X])*(prrl.v[Y] - ppos->v[Y])*(prlr.v[Z] - ppos->v[Z]) / dv;

            lrlq = qdv * (prrl.v[X] - ppos->v[X])*(ppos->v[Y] - plll.v[Y])*(plrr.v[Z] - ppos->v[Z]) / dv;
            rrlq = qdv * (ppos->v[X] - plrl.v[X])*(ppos->v[Y] - prll.v[Y])*(prrr.v[Z] - ppos->v[Z]) / dv;

            rrrq = qdv * (ppos->v[X] - plrr.v[X])*(ppos->v[Y] - prlr.v[Y])*(ppos->v[Z] - prrl.v[Z]) / dv;
            rlrq = qdv * (ppos->v[X] - pllr.v[X])*(prrr.v[Y] - ppos->v[Y])*(ppos->v[Z] - prll.v[Z]) / dv;

            llrq = qdv * (prlr.v[X] - ppos->v[X])*(plrr.v[Y] - ppos->v[Y])*(ppos->v[Z] - plll.v[Z]) / dv;
            lrrq = qdv * (prrr.v[X] - ppos->v[X])*(ppos->v[Y] - pllr.v[Y])*(ppos->v[Z] - plrl.v[Z]) / dv;

            grid->rho[nlc.v[X]][nlc.v[Y]][nlc.v[Z]] += lllq;

            grid->rho[nrc.v[X]][nlc.v[Y]][nlc.v[Z]] += rllq;

            grid->rho[nlc.v[X]][nrc.v[Y]][nlc.v[Z]] += lrlq;

            grid->rho[nlc.v[X]][nlc.v[Y]][nrc.v[Z]] += llrq;

            grid->rho[nrc.v[X]][nrc.v[Y]][nlc.v[Z]] += rrlq;

            grid->rho[nlc.v[X]][nrc.v[Y]][nrc.v[Z]] += lrrq;

            grid->rho[nrc.v[X]][nlc.v[Y]][nrc.v[Z]] += rlrq;

            grid->rho[nrc.v[X]][nrc.v[Y]][nrc.v[Z]] += rrrq;
        }
    }
#endif
}

void ucgrid3d_weight_particles_current(particles_set* set, ucgrid3d* grid)
{
#ifdef VERBOSE
    printf("ucgrid3d_weight_particles_current: weightening\n");
    printf("ucgrid3d_weight_particles_current: using weightening method: %d\n",
           WEIGHTENING_METHOD, BOUNDARIES_TYPE);
#endif
#if ((WEIGHTENING_METHOD == WLINEAR) && (BOUNDARIES_TYPE == CYCLIC_BOUNDARIES))
    speed3d* pspeed;
    pos3d* ppos;
    int i;
    floatdig dv = grid->dx * grid->dy * grid->dz;
    floatdig qdv = set->charge / dv;
    npos3d nlc = {0};
    npos3d nrc = {0};
    pos3d plc = {0};
    pos3d prc = {0};

    pos3d plll = {0};
    pos3d prll = {0};
    pos3d plrl = {0};
    pos3d plrr = {0};
    pos3d pllr = {0};
    pos3d prlr = {0};
    pos3d prrr = {0};
    pos3d prrl = {0};

    floatdig lllq = 0;
    floatdig rllq = 0;
    floatdig lrlq = 0;
    floatdig lrrq = 0;
    floatdig llrq = 0;
    floatdig rlrq = 0;
    floatdig rrrq = 0;
    floatdig rrlq = 0;

    current3d temp = {0};

    for (i = 0; i < set->num_of_particles; i++)
    {
        if (set->particles[i].active == PARTICLE_ACTIVE)
        {
            ppos = &(set->particles[i].pos);
            pspeed = &(set->particles[i].speed);
            ucgrid3d_get_neighbour_cells(grid, ppos, &nlc, &nrc);
            ucgrid3d_get_cell_pos(grid, &nlc, &plc);
            ucgrid3d_get_cell_pos(grid, &nrc, &prc);

            plll.v[X] = plc.v[X];
            plll.v[Y] = plc.v[Y];
            plll.v[Z] = plc.v[Z];

            prll.v[X] = prc.v[X];
            prll.v[Y] = plc.v[Y];
            prll.v[Z] = plc.v[Z];

            plrl.v[X] = plc.v[X];
            plrl.v[Y] = prc.v[Y];
            plrl.v[Z] = plc.v[Z];

            prrl.v[X] = prc.v[X];
            prrl.v[Y] = prc.v[Y];
            prrl.v[Z] = plc.v[Z];

            prrr.v[X] = prc.v[X];
            prrr.v[Y] = prc.v[Y];
            prrr.v[Z] = prc.v[Z];

            prlr.v[X] = prc.v[X];
            prlr.v[Y] = plc.v[Y];
            prlr.v[Z] = prc.v[Z];

            pllr.v[X] = plc.v[X];
            pllr.v[Y] = plc.v[Y];
            pllr.v[Z] = prc.v[Z];

            plrr.v[X] = plc.v[X];
            plrr.v[Y] = prc.v[Y];
            plrr.v[Z] = prc.v[Z];

            lllq = qdv * (prll.v[X] - ppos->v[X])*(plrl.v[Y] - ppos->v[Y])*(pllr.v[Z] - ppos->v[Z]) / dv;
            rllq = qdv * (ppos->v[X] - plll.v[X])*(prrl.v[Y] - ppos->v[Y])*(prlr.v[Z] - ppos->v[Z]) / dv;

            lrlq = qdv * (prrl.v[X] - ppos->v[X])*(ppos->v[Y] - plll.v[Y])*(plrr.v[Z] - ppos->v[Z]) / dv;
            rrlq = qdv * (ppos->v[X] - plrl.v[X])*(ppos->v[Y] - prll.v[Y])*(prrr.v[Z] - ppos->v[Z]) / dv;

            rrrq = qdv * (ppos->v[X] - plrr.v[X])*(ppos->v[Y] - prlr.v[Y])*(ppos->v[Z] - prrl.v[Z]) / dv;
            rlrq = qdv * (ppos->v[X] - pllr.v[X])*(prrr.v[Y] - ppos->v[Y])*(ppos->v[Z] - prll.v[Z]) / dv;

            llrq = qdv * (prlr.v[X] - ppos->v[X])*(plrr.v[Y] - ppos->v[Y])*(ppos->v[Z] - plll.v[Z]) / dv;
            lrrq = qdv * (prrr.v[X] - ppos->v[X])*(ppos->v[Y] - pllr.v[Y])*(ppos->v[Z] - plrl.v[Z]) / dv;

            lib_multiply_vector(pspeed, lllq, &temp);
            lib_add_vector(&grid->j[nlc.v[X]][nlc.v[Y]][nlc.v[Z]], &temp, &grid->j[nlc.v[X]][nlc.v[Y]][nlc.v[Z]]);

            lib_multiply_vector(pspeed, rllq, &temp);
            lib_add_vector(&grid->j[nrc.v[X]][nlc.v[Y]][nlc.v[Z]], &temp, &grid->j[nrc.v[X]][nlc.v[Y]][nlc.v[Z]]);

            lib_multiply_vector(pspeed, lrlq, &temp);
            lib_add_vector(&grid->j[nlc.v[X]][nrc.v[Y]][nlc.v[Z]], &temp, &grid->j[nlc.v[X]][nrc.v[Y]][nlc.v[Z]]);

            lib_multiply_vector(pspeed, llrq, &temp);
            lib_add_vector(&grid->j[nlc.v[X]][nlc.v[Y]][nrc.v[Z]], &temp, &grid->j[nlc.v[X]][nlc.v[Y]][nrc.v[Z]]);

            lib_multiply_vector(pspeed, rrlq, &temp);
            lib_add_vector(&grid->j[nrc.v[X]][nrc.v[Y]][nlc.v[Z]], &temp, &grid->j[nrc.v[X]][nrc.v[Y]][nlc.v[Z]]);

            lib_multiply_vector(pspeed, lrrq, &temp);
            lib_add_vector(&grid->j[nlc.v[X]][nrc.v[Y]][nrc.v[Z]], &temp, &grid->j[nlc.v[X]][nrc.v[Y]][nrc.v[Z]]);

            lib_multiply_vector(pspeed, rlrq, &temp);
            lib_add_vector(&grid->j[nrc.v[X]][nlc.v[Y]][nrc.v[Z]], &temp, &grid->j[nrc.v[X]][nlc.v[Y]][nrc.v[Z]]);

            lib_multiply_vector(pspeed, rrrq, &temp);
            lib_add_vector(&grid->j[nrc.v[X]][nrc.v[Y]][nrc.v[Z]], &temp, &grid->j[nrc.v[X]][nrc.v[Y]][nrc.v[Z]]);
        }
    }
    
#endif

}

void ucgrid3d_dump_rho_into_file(ucgrid3d* grid, char* homedir, char* filename)
{
    int i = 0;
    int j = 0;
    int k = 0;
    char fullname[MAX_BUF_SIZE];
    FILE* to;
    sprintf(fullname, "%s/%s", homedir, filename);
    to = fopen(fullname, "w");
    fprintf(to, "%d\t%d\t%d\n", grid->Nx, grid->Ny, grid->Nz);
    for (i = 0; i < grid->Nx; i++)
    {
        for (j = 0; j < grid->Ny; j++)
        {
            for (k = 0; k < grid->Nz; k++)
            {
                fprintf(to, "%e\n", grid->rho[i][j][k]);
            }
        }
    }
    fclose(to);
}

void ucgrid3d_dump_current_into_file(ucgrid3d* grid, char* homedir, char* filename)
{
    int i = 0;
    int j = 0;
    int k = 0;
    char fullname[MAX_BUF_SIZE];
    FILE* to;
    sprintf(fullname, "%s/%s", homedir, filename);
    to = fopen(fullname, "w");
    fprintf(to, "%d\t%d\t%d\n", grid->Nx, grid->Ny, grid->Nz);
    for (i = 0; i < grid->Nx; i++)
    {
        for (j = 0; j < grid->Ny; j++)
        {
            for (k = 0; k < grid->Nz; k++)
            {
                fprintf(to, "%e\t%e\t%e\n", grid->j[i][j][k].v[X], grid->j[i][j][k].v[Y], grid->j[i][j][k].v[Z]);
            }
        }
    }
    fclose(to);
}
